package imageStitcherAndSplitter;

import java.awt.Graphics;
import java.awt.image.BufferedImage;
import java.awt.image.Raster;
import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Vector;
import java.util.Arrays;

import javax.imageio.ImageIO;

/**
 * Stitches a series of overlapping pictures together to one single picture.
 * Further provides functionality to split the resulting picture into
 * equal size small pictures.
 * 
 * Use case: taking screenshots of google maps in webbrowser to produce 
 * one large image (with very good resolution). Then printing this big picture
 * in a photo shop (like cvs) on multiple small pictures.
 * 
 * @author Nico Zazworka (nico -at- zazworka -dot- com)
 *
 */
public class ImageStitcherAndSplitter {

	
	
	public static void main(String args[]) throws IOException{
		BufferedImage result =doStitching("savedBig8.bmp");
		//doSplitting("saved.bmp");
		
	}
	
	public static void doSplitting(String resultName) throws IOException{
		
		BufferedImage result = ImageIO.read(new File(resultName));
		
		int xRes = 660;
		int yRes = 440;
		
		int xBorder = 66;
		int yBorder= 44;
		
		int xCount = 6;
		int yCount = 6;
		
		for(int i = 0; i<xCount; i++)
			for(int j=0; j<yCount; j++){
				BufferedImage temp = new BufferedImage(xRes+xBorder,yRes+yBorder,BufferedImage.TYPE_INT_RGB);
				Graphics g = temp.getGraphics();
				g.drawImage(result.getSubimage(i*(xRes), j*(yRes), xRes, yRes),xBorder/2,yBorder/2,null);
				//save final image
				try {
				    	System.out.println("writing file output/final_tile_"+i+"_"+j+".png");    
						File outputfile = new File("output/final_tile_"+i+"_"+j+".png");
				        ImageIO.write(temp, "png", outputfile);
				    } catch (IOException e){
				    e.printStackTrace();
				}
			}
		
	}
	
	public static BufferedImage doStitching(String nameFinalImage){
		
		LinkedList<String> files = new LinkedList<String>();
		
		File dir = new File("big");
	    
	    String[] children = dir.list();
	    if (children == null) {
	        // Either dir does not exist or is not a directory
	    } else {
	        for (int i=0; i<children.length  && i<320; i++) {
	            if(children[i].endsWith(".png"))
	            	files.add(children[i]);
	        }
	    }

		System.out.println("Number of images: "+files.size());
	    
		LinkedList<BufferedImage> images = new LinkedList<BufferedImage>();
		
		//load images
		for(int i=0;i<files.size();i++){
			try {
				System.out.println("collegePark/"+files.get(i));
				images.add(ImageIO.read(new File("big/"+files.get(i))));
			} 
			catch (IOException e) {
				e.printStackTrace();
			}
		}
		
		//create stitching matrix:
		// a matrix of stitching points with 
		// matrix[i][j]:stitching point image i to j (null if no stitching found) 
		
		//main set of successfully stitched images 
		List<Integer> mainSet = new LinkedList<Integer>();
		mainSet.add(0);
		
		StitchingPoint[][] sMatrix = new StitchingPoint[images.size()][images.size()];
		int tries=1;
		while(mainSet.size()<images.size() && tries<images.size()){
			System.out.println("\n==== TRIES:"+tries+" ====");
			for(int i=0; i<images.size();i++){
				for(int d=1; d<=tries;d++){
					if(!(!mainSet.contains(i) && (!mainSet.contains(i-d))) &&
							!(mainSet.contains(i) && (mainSet.contains(i-d)))){
						if(i-d>=0){
							int j = i-d;
							System.out.println("Start search for stitching: "+i+(mainSet.contains(i))+"/"+j+(mainSet.contains(j)));
							sMatrix[i][j] = findStitchingPoint(images.get(i), images.get(j));
							
							//stitching found
							if(sMatrix[i][j]!=null){
								System.out.println("  Stitching found : Image "+i+" / "+j+" : "+sMatrix[i][j].x+","+sMatrix[i][j].y+","+sMatrix[i][j].corner );
								if(mainSet.contains(i)){
									mainSet.add(j);
									break;
								}
								else if(mainSet.contains(j)){
									mainSet.add(i);
									break;
								}
							}
						}
					}
					
					if(!(!mainSet.contains(i) && (!mainSet.contains(i+d))) &&
							!(mainSet.contains(i) && (mainSet.contains(i+d)))){	
						if(i+d<images.size()){
							int j = i+d;
							System.out.println("Start search for stitching: "+i+(mainSet.contains(i))+"/"+j+(mainSet.contains(j)));
							sMatrix[i][j] = findStitchingPoint(images.get(i), images.get(j));
							//stitching found
							if(sMatrix[i][j]!=null){
								System.out.println("  Stitching found : Image "+i+" / "+j+" : "+sMatrix[i][j].x+","+sMatrix[i][j].y+","+sMatrix[i][j].corner );
								if(mainSet.contains(i)){
									mainSet.add(j);
									break;
								}
								else if(mainSet.contains(j)){
									mainSet.add(i);
									break;
								}
							}
						}
					}
				}
			
			}
			tries++;
		}
		
		
		//create resulting coordinates for each picture:
		HashMap<Integer,Integer[]> imagesCoordinates = new HashMap<Integer,Integer[]>();
		
		//put first picure in with coordinates (0,0)
		imagesCoordinates.put(0, new Integer[]{0,0});
		
		//iterate and search for connections
		// do as much passes as needed to assign every picture a coordinate
		// at this point it might end up in endless loop if there is a set of non connected pictures
		tries = 0;
		while(imagesCoordinates.size()<images.size() && tries < 100){
			tries++;
			for(int i=0; i<images.size();i++){
				for(int j=0;j<images.size();j++){
					//if
					// - stitching point found between i and j AND
					// - exactly one coordinate of the two is known (XOR)
					if(sMatrix[i][j] != null 
							&& !(imagesCoordinates.containsKey(i) && imagesCoordinates.containsKey(j))
							&& (imagesCoordinates.containsKey(i) || imagesCoordinates.containsKey(j))){
						// i coordinate known -> determine j coordinate
						if(imagesCoordinates.containsKey(i)){
							switch(sMatrix[i][j].corner){
							 	case StitchingPoint.NORTH_WEST:
							 		System.out.println(1);
							 		imagesCoordinates.put(j, new Integer[]
							 		        {imagesCoordinates.get(i)[0]+sMatrix[i][j].x,
							 		         imagesCoordinates.get(i)[1]+sMatrix[i][j].y});
							 		break;
							 	case StitchingPoint.NORTH_EAST:
							 		System.out.println(2);
							 		imagesCoordinates.put(j, new Integer[]
							 		        {imagesCoordinates.get(i)[0]-images.get(j).getWidth()+sMatrix[i][j].x+1,
							 		         imagesCoordinates.get(i)[1]+sMatrix[i][j].y});
							 		break;
							 	case StitchingPoint.SOUTH_EAST:
							 		System.out.println(3);
							 		imagesCoordinates.put(j, new Integer[]
							 		        {imagesCoordinates.get(i)[0]-images.get(j).getWidth()+sMatrix[i][j].x+1,
							 				 imagesCoordinates.get(i)[1]-images.get(j).getHeight()+sMatrix[i][j].y});
							 		break;
							 	case StitchingPoint.SOUTH_WEST:
							 		System.out.println(4);
							 		imagesCoordinates.put(j, new Integer[]
							 		        {imagesCoordinates.get(i)[0]+sMatrix[i][j].x,
							 				imagesCoordinates.get(i)[1]-images.get(j).getHeight()+sMatrix[i][j].y});
							 		break;
							}
							
						}
						else{
							switch(sMatrix[i][j].corner){
						 	case StitchingPoint.NORTH_WEST:
						 		System.out.println(5);
						 		imagesCoordinates.put(i, new Integer[]
						 		        {imagesCoordinates.get(j)[0]-sMatrix[i][j].x,
						 		         imagesCoordinates.get(j)[1]-sMatrix[i][j].y});
						 		break;
						 	case StitchingPoint.NORTH_EAST:
						 		System.out.println(6);
						 		imagesCoordinates.put(i, new Integer[]
						 		        {imagesCoordinates.get(j)[0]+images.get(j).getWidth()-sMatrix[i][j].x-1,
						 		         imagesCoordinates.get(j)[1]-sMatrix[i][j].y});
						 		break;
						 	case StitchingPoint.SOUTH_EAST:
						 		System.out.println(7);
						 		imagesCoordinates.put(i, new Integer[]
						 		        {imagesCoordinates.get(j)[0]+images.get(j).getWidth()-sMatrix[i][j].x-1,
						 				 imagesCoordinates.get(j)[1]+images.get(j).getHeight()-sMatrix[i][j].y});
						 		break;
						 	case StitchingPoint.SOUTH_WEST:
						 		System.out.println(8);
						 		imagesCoordinates.put(i, new Integer[]
						 		        {imagesCoordinates.get(j)[0]-sMatrix[i][j].x,
						 				imagesCoordinates.get(j)[1]+images.get(j).getHeight()-sMatrix[i][j].y});
						 		break;
							}
						}
					}
				}
			}
		}
		
		
		//normalize to (0,0) coordinate
		// - search for x_min and y_min substract from every coordinate
		int x_min = 0 , y_min = 0, x_max = 0, y_max = 0;
		for(int i = 0; i < images.size(); i++){
			Integer[] temp = imagesCoordinates.get(i);
		
			System.out.println(i+" "+temp[0]+"/"+temp[1]);
			
			if(temp[0]<x_min) x_min=temp[0];
			if(temp[0]+images.get(i).getWidth()>x_max) x_max=temp[0]+images.get(i).getWidth();
			if(temp[1]<y_min) y_min=temp[1];
			if(temp[1]+images.get(i).getHeight()>y_max) y_max=temp[1]+images.get(i).getHeight();
		}
		for(Integer[] temp : imagesCoordinates.values()){
			temp[0]=temp[0]-x_min;
			temp[1]=temp[1]-y_min;
		}
		x_max -=x_min;
		y_max -=y_min;
		
		
		//create result image
		BufferedImage finalImage = new BufferedImage(
				x_max,
				y_max,
				BufferedImage.TYPE_INT_RGB);
		
				

		for(int i=0;i<images.size();i++){
			if(imagesCoordinates.containsKey(i)){
				System.out.println("  "+imagesCoordinates.get(i)[0]+"/"+imagesCoordinates.get(i)[1]);
				Graphics g = finalImage.getGraphics();
				g.drawImage(images.get(i),imagesCoordinates.get(i)[0],imagesCoordinates.get(i)[1],null);
			}
		}
		
		
		//save final image
		try {
		        File outputfile = new File(nameFinalImage);
		        ImageIO.write(finalImage, "bmp", outputfile);
		    } catch (IOException e){
		    e.printStackTrace();
		}

		return finalImage;
		
	}

	/**
	 * Tries to map needle to haystack.
	 * @param heystack
	 * @param needle
	 * @return null if no stitching point can be found
	 */
	private static StitchingPoint findStitchingPoint(BufferedImage heystack, BufferedImage needle) {
		
		//10 vertical pixel have to fit
		int stitchingAccuracy = 20;
		
		//for performance
		Raster heystackData = heystack.getData();
		Raster needleData = needle.getData();
		
		for(int x=0;x<heystack.getWidth();x++){
//			if(x%10==0){
//				System.out.println(x);
//				System.out.flush();
//			}
			for(int y=0;y<heystack.getHeight();y++){
				int corner;
				//number of pixels that match
				int fitcount = 0;
				
				//upper left corner
				corner = StitchingPoint.NORTH_WEST;
				for(int i =0 ; i<stitchingAccuracy && y+i<heystack.getHeight();i++){
					//int[] heystackPixel = heystackData.getPixel(x, y+i,new int[]{1,2,3,4});
					
					int[] heystackPixel = heystackData.getPixel(x, y+i,new int[]{1,2,3,4});
					int[] needlePixel = needleData.getPixel(0, i,new int[]{1,2,3,4});
 					if(Arrays.equals(heystackPixel, needlePixel) ){
						fitcount++;
						continue;
					}
 					else{
 						break;
 					}
				}
				if(fitcount>=stitchingAccuracy){
					StitchingPoint result = new StitchingPoint();
					result.x = x;
					result.y = y;
					result.corner = corner;
					return result;
				}
				
				fitcount = 0;
				//upper right corner
				corner = StitchingPoint.NORTH_EAST;
				for(int i =0 ; i<stitchingAccuracy && y+i<heystack.getHeight();i++){
					int[] heystackPixel = heystackData.getPixel(x, y+i,new int[]{1,2,3,4});
					int[] needlePixel = needleData.getPixel(needle.getWidth()-1, i,new int[]{1,2,3,4});
					if(Arrays.equals(heystackPixel, needlePixel) ){
						fitcount++;
						continue;
					}
 					else{
 						break;
 					}
				}
				if(fitcount>=stitchingAccuracy){
					StitchingPoint result = new StitchingPoint();
					result.x = x;
					result.y = y;
					result.corner = corner;
					return result;
				}
				
				fitcount = 0;
				//lower left corner
				corner = StitchingPoint.SOUTH_WEST;
				for(int i =0 ; i<stitchingAccuracy && y+i<heystack.getHeight();i++){	
					int[] heystackPixel = heystackData.getPixel(x, y+i,new int[]{1,2,3,4});
					int[] needlePixel = needleData.getPixel(0, needle.getHeight()-(stitchingAccuracy-i),new int[]{1,2,3,4});
					
					if(Arrays.equals(heystackPixel, needlePixel) ){
						fitcount++;
						continue;
					}
 					else{
 						break;
 					}
				}
				if(fitcount>=stitchingAccuracy){
					StitchingPoint result = new StitchingPoint();
					result.x = x;
					result.y = y+stitchingAccuracy;
					result.corner = corner;
					return result;
				}
				
				fitcount = 0;
				//lower right corner
				corner = StitchingPoint.SOUTH_EAST;
				//for(int i = needle.getHeight()-stitchingAccuracy-1 ; i<needle.getHeight()-1 && y+i<heystack.getHeight();i++){
				for(int i =0 ; i<stitchingAccuracy && y+i<heystack.getHeight();i++){
					int[] heystackPixel = heystackData.getPixel(x, y+i,new int[]{1,2,3,4});
					int[] needlePixel = needleData.getPixel(needle.getWidth()-1, needle.getHeight()-(stitchingAccuracy-i),new int[]{1,2,3,4});
					if(Arrays.equals(heystackPixel, needlePixel) ){
						fitcount++;
						continue;
					}
 					else{
 						break;
 					}
				}
				if(fitcount>=stitchingAccuracy){
					StitchingPoint result = new StitchingPoint();
					result.x = x;
					result.y = y+stitchingAccuracy;
					result.corner = corner;
					return result;
				}
			}
		}
		return null;
	}
}
